import unittest
from alg1 import buildTree, TreeNode

class TestBuildTree(unittest.TestCase):
    def test_example(self):
        preorder = [3, 9, 20, 15, 7]
        inorder = [9, 3, 15, 20, 7]
        root = buildTree(preorder, inorder)
        self.assertEqual(root.val, 3)
        self.assertEqual(root.left.val, 9)
        self.assertEqual(root.right.val, 20)
        self.assertEqual(root.right.left.val, 15)
        self.assertEqual(root.right.right.val, 7)

    def test_single_node(self):
        preorder = [1]
        inorder = [1]
        root = buildTree(preorder, inorder)
        self.assertEqual(root.val, 1)
        self.assertIsNone(root.left)
        self.assertIsNone(root.right)

    def test_left_heavy(self):
        preorder = [1, 2, 3]
        inorder = [3, 2, 1]
        root = buildTree(preorder, inorder)
        self.assertEqual(root.val, 1)
        self.assertEqual(root.left.val, 2)
        self.assertEqual(root.left.left.val, 3)

    def test_right_heavy(self):
        preorder = [1, 2, 3]
        inorder = [1, 2, 3]
        root = buildTree(preorder, inorder)
        self.assertEqual(root.val, 1)
        self.assertEqual(root.right.val, 2)
        self.assertEqual(root.right.right.val, 3)

def tree_to_list(root):
    if root is None:
        return []
    queue = [root]
    result = []
    while queue:
        node = queue.pop(0)
        if node:
            result.append(node.val)
            queue.append(node.left)
            queue.append(node.right)
        else:
            result.append(None)
    # 移除末尾的None
    while result and result[-1] is None:
        result.pop()
    return result

if __name__ == '__main__':
    unittest.main()    